"""
1. 在tensorflow下，训练一个softmax分类器，实现数据集的多分类功能。(100分）
"""
import numpy as np
import tensorflow as tf

np.random.seed(1)
tf.random.set_random_seed(1)

# ①　准备训练数据集x_data、y_data：
x_data = np.array(
    [[1, 2, 3, 4], [2, 1, 1, 2], [3, 1, 2, 4], [4, 2, 5, 5], [1, 7, 5, 5], [1, 4, 5, 6], [1, 6, 9, 6], [1, 7, 9, 7]], dtype=np.float32)
y_data = np.array([[0, 0, 1], [0, 0, 1], [0, 0, 1], [0, 1, 0], [0, 1, 0], [0, 1, 0], [1, 0, 0], [1, 0, 0]], dtype=np.float32)

# ②　定义张量X和Y，float32类型，使用占位符函数。（8分）
X = tf.placeholder(tf.float32, (None, 4), name='X')
Y = tf.placeholder(tf.float32, (None, 3), name='Y')

# ③　定义张量W（weight）和b（bias）。（8分）
W = tf.Variable(tf.random.normal((4, 3)), dtype=tf.float32, name='W')
b = tf.Variable(tf.random.normal((1, 3)), dtype=tf.float32, name='b')

# ④　设置hypothesis预测模型。（8分）
z = tf.matmul(X, W) + b
a = tf.nn.softmax(z)

# ⑤　设置代价函数（损失函数）。（8分）
eps = 1e-20
j = - tf.reduce_mean(tf.matmul(tf.transpose(Y), tf.math.log(a + eps)))
acc = tf.reduce_mean(
    tf.cast(
        tf.equal(
            tf.argmax(a, axis=1),
            tf.argmax(Y, axis=1),
        )
        , dtype=tf.float32
    )
)

# ⑥　使用梯度下降优化器计算最小费用，查找最优解。（8分）
ALPHA = 0.001
dz = a - Y
dW = tf.matmul(tf.transpose(X), dz) / tf.cast(tf.shape(X)[0], dtype=tf.float32)
db = tf.reduce_mean(dz, axis=0)
update = [
    tf.assign(W, W - ALPHA * dW),
    tf.assign(b, b - ALPHA * db),
]

# ⑦　创建会话（Session）。（8分）
with tf.Session() as sess:

    # ⑧　全局变量初始化。（8分）
    sess.run(tf.global_variables_initializer())

    # ⑨　开始迭代总共4001次。（6分）
    ITERS = 4001
    GROUP = 300

    # ⑩　使用训练集的数据进行训练。（8分）
    for i in range(ITERS):
        _, cost, accv = sess.run([update, j, acc], feed_dict={X: x_data, Y: y_data})

        # 11　每300次输出cost值。（6分）
        if i % GROUP == 0:
            print(f'#{i + 1}: cost = {cost}, accv = {accv}')
    if i % GROUP != 0:
        print(f'#{i + 1}: cost = {cost}, accv = {accv}')

    # 12　使用一个测试样本进行分类测试，输出分类结果。（8分）
    # 13　同时使用3个测试样本进行分类测试，输出分类结果。（8分）
    pred = sess.run(a, feed_dict={X: x_data})
    pred_label = np.argmax(pred, axis=1)
    y_label = np.argmax(y_data, axis=1)
    for y_i, pred_label_i in zip(y_label, pred_label):
        print(f'{y_i} => {pred_label_i} ({int(y_i) == pred_label_i})')
